import os
import json
import pandas as pd
from volcenginesdkarkruntime import Ark


def read_local_document(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
        return content
    except FileNotFoundError:
        print(f"错误: 文件 {file_path} 未找到。")
        return None
    except Exception as e:
        print(f"错误: 读取文件时出现问题: {e}")
        return None


def write_to_excel(data, output_file_path):
    try:
        df = pd.DataFrame(data)
        for col in df.columns:
            if df[col].dtype == 'object':
                if col == 'tags':
                    df[col] = df[col].apply(lambda x: ', '.join([tag.replace('[', '').replace(']', '').replace("'", "") for tag in x]) if isinstance(x, list) else x)
                else:
                    df[col] = df[col].apply(lambda x: x.replace('[', '').replace(']', '').replace("'", "") if isinstance(x, str) else x)
        df.to_excel(output_file_path, index=False)
        print(f"已成功将结果写入 {output_file_path}")
    except Exception as e:
        print(f"错误: 写入 Excel 文件时出现问题: {e}")


def get_format_example(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
            example = json.loads(content)
            return example
    except FileNotFoundError:
        print(f"错误: 格式示例文件 {file_path} 未找到。")
        return None
    except json.JSONDecodeError as e:
        print(f"错误: 解析格式示例文件 {file_path} 时出现问题: {e}，文件内容: {content}")
        return None


def write_raw_content_to_json(raw_content, file_path):
    try:
        # 尝试解析原始内容为 JSON 格式
        if isinstance(raw_content, str):
            try:
                parsed_content = json.loads(raw_content)
                # 如果解析成功，按照格式写入文件
                with open(file_path, 'w', encoding='utf-8') as file:
                    json.dump(parsed_content, file, ensure_ascii=False, indent=4)
                print(f"已成功将大模型返回的原始内容写入 {file_path}")
            except json.JSONDecodeError:
                # 如果解析失败，尝试按照每行是一个 JSON 对象的格式处理
                lines = raw_content.splitlines()
                valid_lines = [line for line in lines if line.strip()]
                json_objects = []
                for line in valid_lines:
                    try:
                        obj = json.loads(line)
                        json_objects.append(obj)
                    except json.JSONDecodeError:
                        pass
                with open(file_path, 'w', encoding='utf-8') as file:
                    file.write('[')
                    for index, obj in enumerate(json_objects):
                        json.dump(obj, file, ensure_ascii=False, indent=4)
                        if index < len(json_objects) - 1:
                            file.write(',')
                    file.write(']')
                print(f"已成功将大模型返回的原始内容整理后写入 {file_path}")
    except Exception as e:
        print(f"错误: 写入大模型返回的原始内容时出现问题: {e}")


# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
# 初始化 Ark 客户端，从环境变量中读取您的 API Key
client = Ark(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    api_key=os.environ.get("ARK_API_KEY"),
)

input_file_path = 'test_document.txt'
output_file_path = 'extracted_testpoints.xlsx'
format_example_path = 'case.json'
raw_content_output_path = 'extracted_testpoints.json'
document_content = read_local_document(input_file_path)
format_example = get_format_example(format_example_path)

if document_content and format_example:
    example_str = json.dumps(format_example, ensure_ascii=False)
    response = client.chat.completions.create(
        model="ep-20250321145642-7fkgf",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"请从以下文档内容中提炼测试点，涵盖异常场景、性能场景和安全场景。将结果整理成与下面示例格式一致的 JSON 测试用例。如果存在多个测试场景，请输出对应个数的测试用例。示例格式：{example_str}。文档内容：{document_content}"
                    }
                ]
            }
        ]
    )
    raw_content = response.choices[0].message.content
    print(f"大模型返回的原始内容：{raw_content}")
    write_raw_content_to_json(raw_content, raw_content_output_path)
    try:
        test_cases = json.loads(raw_content)
        write_to_excel(test_cases, output_file_path)
    except json.JSONDecodeError as e:
        print(f"解析 JSON 时出现错误：{e}")
        print("尝试根据格式示例进行调整...")
        try:
            adjusted_content = raw_content.strip()
            if adjusted_content.startswith("```json") and adjusted_content.endswith("```"):
                adjusted_content = adjusted_content[7:-3].strip()
            test_cases = json.loads(adjusted_content)
            write_to_excel(test_cases, output_file_path)
        except json.JSONDecodeError as inner_e:
            print(f"调整后仍无法解析 JSON：{inner_e}")